package com.yww.datastructure.binarysearchtree;

import com.yww.datastructure.line.Linear;
import com.yww.datastructure.line.array.DynamicArray;
import com.yww.datastructure.line.linkedlist.CircularLinkedList;
import com.yww.datastructure.line.linkedlist.DoubleLinkedList;
import com.yww.datastructure.stack.LinkedListStack;
import com.yww.datastructure.stack.Stack;

import java.util.function.Consumer;

/**
 * @author yww
 * @description key为int的二叉搜索树
 * @since 2024/2/14 10:34
 */
public final class IntBinarySearchTree<V> {

    private BSTNode<V> root;

    /**
     * 查询方法
     */
    public V get(int key) {
        return doGet(key);
    }

    /**
     * 查询方法
     */
    public V recurGet(int key) {
        return recurGet(key, root);
    }

    /**
     * 循环查询
     */
    private V doGet(int key) {
        if (root == null) {
            return null;
        }
        BSTNode<V> p = root;
        while (p != null) {
            if (key < p.key) {
                p = p.left; // 向左查找
            } else if (p.key < key) {
                p = p.right; // 向右查找
            } else {
                return p.val;
            }
        }
        return null; // 没找到
    }

    /**
     * 递归查询
     */
    private V recurGet(int key, BSTNode<V> start) {
        if (start == null) {
            return null; // 没找到
        }
        if (key < start.key) {
            return recurGet(key, start.left); // 向左找
        } else if (start.key < key) {
            return recurGet(key, start.right); // 向右找
        } else {
            return start.val; // 找到了
        }
    }

    /**
     * 循环查最小
     */
    public V min() {
        if (root == null) {
            return null;
        }
        BSTNode<V> p = root;
        while (p.left != null) {
            p = p.left;
        }
        return p.val;
    }

    /**
     * 递归查最小
     */
    public V recurMin() {
        return recurMin(root);
    }

    private V recurMin(BSTNode<V> start) {
        if (start == null) {
            return null;
        }
        if (start.left == null) {
            return start.val;
        }
        return recurMin(start.left);
    }

    /**
     * 循环查最大
     */
    public V max() {
        if (root == null) {
            return null;
        }
        BSTNode<V> p = root;
        while (p.right != null) {
            p = p.right;
        }
        return p.val;
    }

    /**
     * 递归查最大
     */
    public V recurMax() {
        return recurMax(root);
    }

    private V recurMax(BSTNode<V> start) {
        if (start == null) {
            return null;
        }
        if (start.right == null) {
            return start.val;
        }
        return recurMax(start.right);
    }

    /**
     * 新增
     */
    public void put(int key, V val) {
        BSTNode<V> p = root;
        BSTNode<V> parent = null;
        while (p != null) {
            parent = p; // 记录一下父节点，为新增做准备
            if (key < p.key) {
                p = p.left;
            } else if (p.key < key) {
                p = p.right;
            } else { // 找到相同的key，val更新
                p.val = val;
                return;
            }
        }
        // key不存在，则新增
        if (parent == null) {
            root = new BSTNode<>(key, val);
        } else if (key < parent.key) {
            parent.left = new BSTNode<>(key, val);
        } else {
            parent.right = new BSTNode<>(key, val);
        }

    }

    /**
     * 递归新增
     */
    public void recurPut(int key, V val) {
        recurPut(key, val, root);
    }

    private BSTNode<V> recurPut(int key, V val, BSTNode<V> start) {
        if (start == null) {
            return new BSTNode<>(key, val); // 创建新节点
        }
        if (key < start.key) {
            start.left = recurPut(key, val, start.left); // 左孩子为左递归创建的新节点
        } else if (start.key < key) {
            start.right = recurPut(key, val, start.right); // 右孩子为右递归创建的新节点
        } else {
            start.val = val; // 存在相同的key，更新
        }
        return start;
    }

    /**
     * 找key的前驱 <br>
     * 一个节点的前驱（前任）节点是指比它小的节点中，最大的那个 <br>
     * 节点有左子树，此时前驱节点就是左子树的最大值 <br>
     * 节点没有左子树，若离它最近的祖先自从左而来，此祖先即为前驱 <br>
     */
    public V predecessor(int key) {
        BSTNode<V> p = root;
        BSTNode<V> ancestorFromLeft = null;
        while (p != null) { // 找到key节点
            if (key < p.key) {
                p = p.left;
            } else if (p.key < key) {
                ancestorFromLeft = p; // 记录自左而来的祖先节点
                p = p.right;
            } else {
                break;
            }
        }
        if (p == null) {
            return null;
        }
        if (p.left != null) {
            return recurMax(p.left); // 有左子树，则前驱为左子树最大值
        }
        return ancestorFromLeft == null ? null : ancestorFromLeft.val;
    }

    /**
     * 找key的后继 <br>
     * 一个节点的后继（后任）节点是指比它大的节点中，最小的那个 <br>
     * 节点有右子树，此时后继节点就是右子树的最小值 <br>
     * 节点没有右子树，若离它最近的祖先自从右而来，此祖先即为后继
     */
    public V successor(int key) {
        BSTNode<V> p = root;
        BSTNode<V> ancestorFromRight = null;
        while (p != null) { // 找到为key的节点
            if (key < p.key) {
                ancestorFromRight = p; // 记录自右而来的祖先节点
                p = p.left;
            } else if (p.key < key) {
                p = p.right;
            } else {
                break;
            }
        }
        if (p == null) {
            return null;
        }
        if (p.right != null) { // 有右子树，则后继为右子树的最小值
            return recurMin(p.right);
        }
        return ancestorFromRight == null ? null : ancestorFromRight.val;
    }

    /**
     * 非递归删除方法 <br>
     * 要删除某节点（称为 D），必须先找到被删除节点的父节点，这里称为 Parent <br>
     * <ul>
     *     <li>1.删除节点没有左孩子，将右孩子托孤给 Parent</li>
     *     <li>2.删除节点没有右孩子，将左孩子托孤给 Parent</li>
     *     <li>3.删除节点左右孩子都没有，已经被涵盖在情况1、情况2 当中，把 null 托孤给 Parent</li>
     *     <li>
     *        4.删除节点左右孩子都有，可以将它的后继节点（称为 S）托孤给 Parent，设 S 的父亲为 SP，又分两种情况  <br>
     *        <ul>
     *              <li>(1) SP 就是被删除节点，此时 D 与 S 紧邻，只需将 S 托孤给 Parent</li>
     *              <li>(2) SP 不是被删除节点，此时 D 与 S 不相邻，此时需要将 S 的后代托孤给 SP，再将 S 托孤给 Parent</li>
     *        </ul>
     *     </li>
     * </ul>
     */
    public V delete(int key) {
        BSTNode<V> p = root;
        BSTNode<V> parent = null;
        while (p != null) { // 找到待删除的p节点
            if (key < p.key) {
                parent = p; // 记录父节点
                p = p.left;
            } else if (p.key < key) {
                parent = p; // 记录父节点
                p = p.right;
            } else {
                break;
            }
        }
        if (p == null) {
            return null;
        }
        if (p.left == null) { // 情况1，没左孩子，右孩子托孤给父，情况3包含
            shift(parent, p, p.right);
        } else if (p.right == null) { // 情况2，没右孩子，左孩子托孤给父，情况3包含
            shift(parent, p, p.left);
        } else { // 情况4，删除节点有左右孩子
            BSTNode<V> s = p.right; // 找到删除节点的后继节点
            BSTNode<V> sp = p; // 后继节点的父节点
            while (s.left != null) {
                sp = s;
                s = s.left;
            }

            if (sp != p) { // 如果删除节点与后继节点不紧密相连，要处理后继节点的后事
                shift(sp, s, s.right); // 不可能右左孩子
                s.right = p.right; // 将待删除节点的右孩子赋值给后继节点
            }
            shift(parent, p, s); // 将后继节点顶上去
            s.left = p.left; // 将待删除节点的左孩子赋值给后继节点
        }
        return p.val;
    }

    /**
     * 托孤方法 <br>
     *
     * @param parent  被删除节点的父节点
     * @param deleted 被删除节点
     * @param child   被顶上去的节点
     */
    private void shift(BSTNode<V> parent, BSTNode<V> deleted, BSTNode<V> child) {
        if (parent == null) {
            root = child;
        } else if (deleted == parent.left) {
            parent.left = child;
        } else {
            parent.right = child;
        }
    }

    /**
     * 递归删除
     */
    public V recurDel(int key) {
        Linear<V> delRes = new DoubleLinkedList<>();
        root = recurDel(root, key, delRes);
        return delRes.isEmpty() ? null : delRes.get(0);
    }

    /**
     * 递归删除
     *
     * @param start  从那个节点开始删除
     * @param key    要删除节点的key
     * @param delRes 保存删除节点的val，该集合只会保留一个元素，如果找到删除的节点
     * @return 成功删除节点之后剩余的节点，null没找到
     */
    private BSTNode<V> recurDel(BSTNode<V> start, int key, Linear<V> delRes) {
        if (start == null) {
            return null; // 没找到
        }
        if (key < start.key) {
            start.left = recurDel(start.left, key, delRes);
            return start; // 一定要返回重新建立关系的当前节点
        }
        if (start.key < key) {
            start.right = recurDel(start.right, key, delRes);
            return start; // 一定要返回重新建立关系的当前节点
        }
        // 找到了，保存待删除节点的val
        delRes.addLast(start.val);
        // 情况3：左右孩子都有
        if (start.left != null && start.right != null) {
            BSTNode<V> s = start.right; // 找到删除节点的后继节点
            while (s.left != null) {
                s = s.left;
            }
            s.right = recurDel(start.right, s.key, new CircularLinkedList<>()); // 更新后继节点的右指针（先）
            s.left = start.left; // 更新后继节点的左指针（后）
            return s; // 返回更新完成的后继节点
        }
        // 情况1：只有左孩子 情况2：只有右孩子
        return start.left != null ? start.left : start.right;
    }

    public void inorderTraversal(Consumer<V> consumer) {
        BSTNode<V> cur = root;
        Stack<BSTNode<V>> stack = new LinkedListStack<>();
        BSTNode<V> pop = null;
        while (!stack.isEmpty() || cur != null) {
            if (cur != null) { // 当前节点不空
                stack.push(cur);
                cur = cur.left; // 访问左
            } else {
                BSTNode<V> peek = stack.peek(); // 回
                if (peek.right == null) { // 没有右子树，无需处理，直接弹出
                    if (consumer != null) {
                        consumer.accept(peek.val);
                    }
                    pop = stack.pop();
                } else if (peek.right == pop) { // 右子树已经处理，直接弹出
                    pop = stack.pop();
                } else { // 右子树未处理，需要处理
                    if (consumer != null) {
                        consumer.accept(peek.val);
                    }
                    cur = peek.right;
                }
            }
        }
    }

    /**
     * 返回比key小的元素集合
     */
    public Linear<V> less(int key) {
        Linear<V> lessRes = new DynamicArray<>();
        BSTNode<V> cur = root;
        Stack<BSTNode<V>> stack = new LinkedListStack<>();
        while (!stack.isEmpty() || cur != null) {
            if (cur != null) {
                stack.push(cur);
                cur = cur.left;
            } else {
                BSTNode<V> pop = stack.pop();
                if (pop.key < key) {
                    lessRes.addLast(pop.val);
                } else {
                    break;
                }
                cur = pop.right;
            }
        }
        return lessRes;
    }

    /**
     * 返回在key1 ~ key2之间的元素集合
     */
    public Linear<V> between(int key1, int key2) {
        Linear<V> betweenRes = new DynamicArray<>();
        BSTNode<V> cur = root;
        Stack<BSTNode<V>> stack = new LinkedListStack<>();
        while (!stack.isEmpty() || cur != null) {
            if (cur != null) {
                stack.push(cur);
                cur = cur.left;
            } else {
                BSTNode<V> pop = stack.pop();
                if (pop.key >= key1 && pop.key <= key2) {
                    betweenRes.addLast(pop.val);
                } else if (pop.key > key2) {
                    break;
                }
                cur = pop.right;
            }
        }
        return betweenRes;
    }

    /**
     * 返回比key大的元素集合
     * 使用了逆中序遍历 RNL
     */
    public Linear<V> greater(int key) {
        Linear<V> greaterRes = new DynamicArray<>();
        BSTNode<V> cur = root;
        Stack<BSTNode<V>> stack = new LinkedListStack<>();
        while (!stack.isEmpty() || cur != null) {
            if (cur != null) {
                stack.push(cur);
                cur = cur.right;
            } else {
                BSTNode<V> pop = stack.pop();
                if (pop.key > key) {
                    greaterRes.addFirst(pop.val);
                } else {
                    break;
                }
                cur = pop.left;
            }
        }
        return greaterRes;
    }


    /**
     * 树节点
     */
    private static class BSTNode<V> {
        private final int key;
        private V val;
        private BSTNode<V> left;
        private BSTNode<V> right;

        public BSTNode(int key, V val) {
            this.key = key;
            this.val = val;
        }

        public BSTNode(int key, V val, BSTNode<V> left, BSTNode<V> right) {
            this.key = key;
            this.val = val;
            this.left = left;
            this.right = right;
        }

        @Override
        public String toString() {
            return "{" +
                    "key=" + key +
                    ", val=" + val +
                    '}';
        }
    }
}
